import argparse
from typing import List

import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg

from dataset import get_transforms
from model.AllModel import Encoder, Decoder
from chemistry import convert_graph_to_smiles
from tokenizer import get_tokenizer

# 定义了不同类型的化学键
BOND_TYPES = ["", "single", "double", "triple", "aromatic", "solid wedge", "dashed wedge"]


def safe_load(module, module_states):
    def remove_prefix(state_dict):
        return {k.replace('module.', ''): v for k, v in state_dict.items()}
    missing_keys, unexpected_keys = module.load_state_dict(remove_prefix(module_states), strict=False)
    return


class MolScribe:

    def __init__(self, model_path, device=None, num_workers=1):
        model_states = torch.load(model_path, map_location=torch.device('cpu'))
        args = self._get_args(model_states['args'])
        if device is None:
            device = torch.device('cpu')
        self.device = device
        self.tokenizer = get_tokenizer(args)
        self.encoder, self.decoder = self._get_model(args, self.tokenizer, self.device, model_states)
        self.transform = get_transforms(args.input_size, augment=False)
        self.num_workers = num_workers

    # 使用 Python 的 argparse 库来解析命令行参数，并构建了一个参数解析器 parser。
    # 这个方法允许你定义一系列命令行参数，这些参数可以用于配置模型、数据和训练过程
    def _get_args(self, args_states=None):

        # 创建参数解析器
        parser = argparse.ArgumentParser()
        # Model
        parser.add_argument('--encoder', type=str, default='swin_base')
        parser.add_argument('--decoder', type=str, default='transformer')
        parser.add_argument('--trunc_encoder', action='store_true')  # use the hidden states before downsample
        parser.add_argument('--no_pretrained', action='store_true')
        parser.add_argument('--use_checkpoint', action='store_true', default=True)
        parser.add_argument('--dropout', type=float, default=0.5)
        parser.add_argument('--embed_dim', type=int, default=256)
        parser.add_argument('--enc_pos_emb', action='store_true')

        # 专门用于 Transformer 模型的配置参数
        group = parser.add_argument_group("transformer_options")
        group.add_argument("--dec_num_layers", help="No. of layers in transformer decoder", type=int, default=6)
        group.add_argument("--dec_hidden_size", help="Decoder hidden size", type=int, default=256)
        group.add_argument("--dec_attn_heads", help="Decoder no. of attention heads", type=int, default=8)
        group.add_argument("--dec_num_queries", type=int, default=128)
        group.add_argument("--hidden_dropout", help="Hidden dropout", type=float, default=0.1)
        group.add_argument("--attn_dropout", help="Attention dropout", type=float, default=0.1)
        group.add_argument("--max_relative_positions", help="Max relative positions", type=int, default=0)
        parser.add_argument('--continuous_coords', action='store_true')
        parser.add_argument('--compute_confidence', action='store_true')

        # Data,添加了与数据相关的参数，如输入大小、词汇表文件、坐标箱数等
        parser.add_argument('--input_size', type=int, default=384)
        parser.add_argument('--vocab_file', type=str, default=None)
        parser.add_argument('--coord_bins', type=int, default=64)
        parser.add_argument('--sep_xy', action='store_true', default=True)

        # 解析命令行参数。由于传入了一个空列表 []，这意味着不会解析实际的命令行参数，而是使用默认值
        args = parser.parse_args([])

        # 如果提供了 args_states 字典，那么遍历该字典，并更新 args 对象的属性。
        # 这允许在不通过命令行参数的情况下，从代码中设置参数值
        if args_states:
            for key, value in args_states.items():
                args.__dict__[key] = value
        return args

    def _get_model(self, args, tokenizer, device, states):
        encoder = Encoder(args, pretrained=False)
        args.encoder_dim = encoder.n_features
        decoder = Decoder(args, tokenizer)

        safe_load(encoder, states['encoder'])
        safe_load(decoder, states['decoder'])
        # print(f"Model loaded from {load_path}")

        encoder.to(device)
        decoder.to(device)
        encoder.eval()
        decoder.eval()
        return encoder, decoder
    
    # 一批图像进行处理，并预测其中的化学分子结构
    # 对输入的图像列表进行批处理，每批处理 batch_size 个图像。
    # 对每批图像应用预处理转换。
    # 使用编码器（Encoder）提取图像特征。
    # 使用解码器（Decoder）根据提取的特征预测化学结构。
    # 将预测的结构转换为SMILES字符串和分子块（Molfile）。
    # 如果需要，返回原子和化学键的详细信息以及置信度。
    

# return_atoms_bonds 是一个布尔类型的参数，它通常在函数或方法中用来指示是否需要返回原子和化学键的详细信息。
# 在 MolScribe 类的上下文中，这个参数决定了在预测化学结构时，
# 除了基本的 SMILES 字符串和分子块（Molfile）之外，是否还要包括原子坐标、原子符号以及原子之间的化学键信息。
    def predict_images(self, input_images: List, return_atoms_bonds=False, return_confidence=False, batch_size=16):
        device = self.device
        predictions = []
        self.decoder.compute_confidence = return_confidence
    # 分好批次
        for idx in range(0, len(input_images), batch_size):
            # 划分图片批次
            batch_images = input_images[idx:idx+batch_size]
            # 图像预处理
            images = [self.transform(image=image, keypoints=[])['image'] for image in batch_images]
            # 图像堆叠，批次处理
            images = torch.stack(images, dim=0).to(device)
            with torch.no_grad():
                features, hiddens = self.encoder(images)
                # 为何这个能直接使用？
                batch_predictions = self.decoder.decode(features, hiddens)
                print("====也进入了decoder.decode了啊？===")
            predictions += batch_predictions
            # 解码器输出结构
                
#             predictions = [
#     {
#         'chartok_coords': {
#             'smiles': 'CCO',
#             'coords': [[1.0, 2.0], [2.0, 3.0], [3.0, 4.0]],
#             'symbols': ['C', 'C', 'O']
#         },
#         'edges': {
#             'bonds': [[0, 1, 1], [1, 2, 1]],  # atom index, atom index, bond type
#             'bond_types': ['single', 'single']
#         },
#         'confidence': 0.95
#     },
#     # ... 更多预测结果
# ]
                

            
        smiles = [pred['chartok_coords']['smiles'] for pred in predictions]

        #  一个列表，包含每个原子的坐标
        node_coords = [pred['chartok_coords']['coords'] for pred in predictions]

        # 一个列表，包含每个原子的符号
        node_symbols = [pred['chartok_coords']['symbols'] for pred in predictions]

        # 一个列表，包含化学键的信息
        edges = [pred['edges'] for pred in predictions]
        
        # 用获得的化学信息，对SMILES进行化学重构，使得比预测得到的smiles更准确
        smiles_list, molblock_list, r_success = convert_graph_to_smiles(
            node_coords, node_symbols, edges, images=input_images, num_workers=self.num_workers)

        outputs = []
        for smiles, molblock, pred in zip(smiles_list, molblock_list, predictions):
            pred_dict = {"smiles": smiles, "molfile": molblock}
            if return_confidence:
                pred_dict["confidence"] = pred["overall_score"]
            if return_atoms_bonds:
                coords = pred['chartok_coords']['coords']
                symbols = pred['chartok_coords']['symbols']
                # get atoms info
                atom_list = []
                for i, (symbol, coord) in enumerate(zip(symbols, coords)):
                    # 元素的坐标，名称信息
                    atom_dict = {"atom_symbol": symbol, "x": coord[0], "y": coord[1]}
                    if return_confidence:
                        atom_dict["confidence"] = pred['chartok_coords']['atom_scores'][i]
                    atom_list.append(atom_dict)
                pred_dict["atoms"] = atom_list
                # get bonds info
                bond_list = []
                num_atoms = len(symbols)
                for i in range(num_atoms-1):
                    for j in range(i+1, num_atoms):
                        # 从解码器的预测中找到对i原子和j原子的键预测
                        bond_type_int = pred['edges'][i][j]
                        if bond_type_int != 0:
                            bond_type_str = BOND_TYPES[bond_type_int]
                            bond_dict = {"bond_type": bond_type_str, "endpoint_atoms": (i, j)}
                            if return_confidence:
                                bond_dict["confidence"] = pred["edge_scores"][i][j]
                            bond_list.append(bond_dict)
                pred_dict["bonds"] = bond_list
            outputs.append(pred_dict)
        return outputs
    # 输出结构
#    [
#       {
#     "smiles": "CCO",
#     "molfile": "...",
#     "confidence": 0.95,
#     "atoms": [
#         {"atom_symbol": "C", "x": 1.0, "y": 2.0, "confidence": 0.9},
#         {"atom_symbol": "C", "x": 2.0, "y": 3.0, "confidence": 0.92},
#         {"atom_symbol": "O", "x": 3.0, "y": 4.0, "confidence": 0.91}
#     ],
#     "bonds": [
#         {"bond_type": "single", "endpoint_atoms": (0, 1), "confidence": 0.89},
#         {"bond_type": "single", "endpoint_atoms": (1, 2), "confidence": 0.9}
#     ]
#       },
#       .......
#    ] 


    #简化版的predict_images, 是对单张图片处理 
    def predict_image(self, image, return_atoms_bonds=False, return_confidence=False):
        # outputs[0]--->{"smiles": "CCO", "molfile": "..."}
        return self.predict_images([
            image], return_atoms_bonds=return_atoms_bonds, return_confidence=return_confidence)[0]

    def predict_image_files(self, image_files: List, return_atoms_bonds=False, return_confidence=False):
        input_images = []
        # 图片的读取，处理，存储
        for path in image_files:
            image = cv2.imread(path)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            input_images.append(image)
        # 然后用predict_images批量预测
        return self.predict_images(
            input_images, return_atoms_bonds=return_atoms_bonds, return_confidence=return_confidence)
    
    # predict_image_files的简化版，单张图片读取预测
    def predict_image_file(self, image_file: str, return_atoms_bonds=False, return_confidence=False):
        return self.predict_image_files(
            [image_file], return_atoms_bonds=return_atoms_bonds, return_confidence=return_confidence)[0]

    def draw_prediction(self, prediction, image, notebook=False):
        if "atoms" not in prediction or "bonds" not in prediction:
            raise ValueError("atoms and bonds information are not provided.")
        h, w, _ = image.shape
        # 将图像的高度和宽度调整到最大边长为 400 像素，同时保持图像的宽高比不变
        h, w = np.array([h, w]) * 400 / max(h, w)
        image = cv2.resize(image, (int(w), int(h)))
        fig, ax = plt.subplots(1, 1)
        ax.axis('off')
        ax.set_xlim(-0.05 * w, w * 1.05)
        ax.set_ylim(1.05 * h, -0.05 * h)
        plt.imshow(image, alpha=0.)
        # 因为对图像进行了尺寸调整，那么对于预测结果中的原子坐标点也需要进行相应的调整
        x = [a['x'] * w for a in prediction['atoms']]
        y = [a['y'] * h for a in prediction['atoms']]
        markersize = min(w, h) / 3
        plt.scatter(x, y, marker='o', s=markersize, color='lightskyblue', zorder=10)
        # 每个原子的位置上方绘制原子符号的循环
        for i, atom in enumerate(prediction['atoms']):
            symbol = atom['atom_symbol'].lstrip('[').rstrip(']')
            plt.annotate(symbol, xy=(x[i], y[i]), ha='center', va='center', color='black', zorder=100)
        for bond in prediction['bonds']:
            u, v = bond['endpoint_atoms']
            # 把有键的两个原子坐标拿出来
            x1, y1, x2, y2 = x[u], y[u], x[v], y[v]
            bond_type = bond['bond_type']
            # 'single'：单键，使用绿色，线宽为 4。
            # 'aromatic'：芳香键，使用紫色，线宽为 4。
            # 'double'：双键，使用绿色，线宽为 7，并且添加一条白色的细线以表示双键的厚度。
            # 'triple'：三键，使用绿色，通过绘制三条线来表示三键，中间的线稍微细一些。
            if bond_type == 'single':
                color = 'tab:green'
                ax.plot([x1, x2], [y1, y2], color, linewidth=4)
            elif bond_type == 'aromatic':
                color = 'tab:purple'
                ax.plot([x1, x2], [y1, y2], color, linewidth=4)
            elif bond_type == 'double':
                color = 'tab:green'
                ax.plot([x1, x2], [y1, y2], color=color, linewidth=7)
                ax.plot([x1, x2], [y1, y2], color='w', linewidth=1.5, zorder=2.1)
            elif bond_type == 'triple':
                color = 'tab:green'
                x1s, x2s = 0.8 * x1 + 0.2 * x2, 0.2 * x1 + 0.8 * x2
                y1s, y2s = 0.8 * y1 + 0.2 * y2, 0.2 * y1 + 0.8 * y2
                ax.plot([x1s, x2s], [y1s, y2s], color=color, linewidth=9)
                ax.plot([x1, x2], [y1, y2], color='w', linewidth=5, zorder=2.05)
                ax.plot([x1, x2], [y1, y2], color=color, linewidth=2, zorder=2.1)
            else:
                length = 10
                width = 10
                color = 'tab:green'
                if bond_type == 'solid wedge':
                    ax.annotate('', xy=(x1, y1), xytext=(x2, y2),
                                arrowprops=dict(color=color, width=3, headwidth=width, headlength=length), zorder=2)
                else:
                    ax.annotate('', xy=(x2, y2), xytext=(x1, y1),
                                arrowprops=dict(color=color, width=3, headwidth=width, headlength=length), zorder=2)
        fig.tight_layout()
        if not notebook:
            canvas = FigureCanvasAgg(fig)
            canvas.draw()
            buf = canvas.buffer_rgba()
            result_image = np.asarray(buf)
            plt.close(fig)
            return result_image
